"""Natural language query engine for code repositories."""

import json
from typing import List, Dict, Any, Optional
from pathlib import Path

from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema import Document
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
from rich.console import Console
from rich.markdown import Markdown

from .config import Config
from .rag_indexer import RAGIndexer

console = Console()

class QueryEngine:
    """Natural language query engine for code repositories."""
    
    def __init__(self):
        self.llm = ChatOpenAI(
            model_name=Config.CHAT_MODEL,
            temperature=0.1,
            openai_api_key=Config.KIMI_API_KEY,
            openai_api_base=Config.KIMI_API_BASE
        )
        self.indexer = RAGIndexer()
        
        # Define query templates
        self._setup_prompts()
    
    def _setup_prompts(self):
        """Setup prompt templates for different query types."""
        
        # General query template
        self.general_prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a helpful code assistant that helps developers understand code repositories.
            
            You have access to code snippets from various files in the repository. When answering questions:
            1. Be specific and provide code examples when relevant
            2. Include file paths and line numbers when possible
            3. Explain complex concepts clearly
            4. If you're unsure about something, say so
            
            Context:
            {context}
            """),
            ("human", "{question}")
        ])
        
        # Function explanation template
        self.function_explanation_prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a code documentation expert. Explain the purpose, usage, and implementation of the given function.
            
            Provide:
            1. A clear explanation of what the function does
            2. Input parameters and their types
            3. Return value and type
            4. Any side effects or important behavior
            5. Usage examples if available
            6. Related functions or classes
            
            Code:
            {code}
            
            Context:
            {context}
            """),
            ("human", "Explain this function: {function_name}")
        ])
        
        # Module usage template
        self.module_usage_prompt = ChatPromptTemplate.from_messages([
            ("system", """You are analyzing how a specific module, library, or component is used in this codebase.
            
            For each usage found, provide:
            1. The file path where it's used
            2. The specific context of usage
            3. How it's being used (imported, called, extended, etc.)
            4. Any patterns or common usage you notice
            
            Context:
            {context}
            """),
            ("human", "How is the module '{module_name}' used in this codebase?")
        ])
        
        # Error handling template
        self.error_handling_prompt = ChatPromptTemplate.from_messages([
            ("system", """You are analyzing error handling patterns in this codebase.
            
            For error handling examples found, provide:
            1. The specific error type being handled
            2. The file path and context
            3. How the error is being handled (try-catch, error returns, etc.)
            4. Any best practices or patterns observed
            
            Context:
            {context}
            """),
            ("human", "Find examples of handling {error_type} errors in this codebase")
        ])
    
    def query_repository(self, query: str, repo_url: str = None) -> str:
        """
        Execute a natural language query against the indexed repository.
        
        Args:
            query: Natural language question
            repo_url: Optional repository URL to filter results
            
        Returns:
            Formatted answer as markdown
        """
        try:
            # First, try to understand the query type
            query_type = self._classify_query(query)
            
            # Get relevant context
            context_docs = self._get_relevant_context(query, query_type, repo_url)
            
            if not context_docs:
                return "I couldn't find any relevant code snippets for your query."
            
            # Generate response based on query type
            response = self._generate_response(query, query_type, context_docs)
            
            return response
            
        except Exception as e:
            console.print(f"[red]Error processing query: {e}[/red]")
            return f"Sorry, I encountered an error: {str(e)}"
    
    def _classify_query(self, query: str) -> str:
        """Classify the type of query."""
        query_lower = query.lower()
        
        if any(word in query_lower for word in ['what does', 'explain', 'describe', 'function']):
            return 'function_explanation'
        elif any(word in query_lower for word in ['import', 'use', 'usage', 'module', 'library']):
            return 'module_usage'
        elif any(word in query_lower for word in ['error', 'exception', 'handling', 'catch']):
            return 'error_handling'
        elif any(word in query_lower for word in ['where', 'find', 'locate']):
            return 'location_search'
        else:
            return 'general'
    
    def _get_relevant_context(self, query: str, query_type: str, repo_url: str = None) -> List[Document]:
        """Get relevant code snippets based on query type."""
        
        if query_type == 'function_explanation':
            # Extract function name from query
            function_name = self._extract_function_name(query)
            if function_name:
                return self.indexer.search_functions(function_name, repo_url)
        
        elif query_type == 'module_usage':
            # Extract module name from query
            module_name = self._extract_module_name(query)
            if module_name:
                return self.indexer.search_imports(module_name, repo_url)
        
        elif query_type == 'error_handling':
            # Search for error handling patterns
            error_type = self._extract_error_type(query)
            search_query = f"error handling {error_type}" if error_type else "error handling"
            return self.indexer.query_similar_chunks(search_query, k=10)
        
        # Default: search with the query
        filter_dict = {"repo_url": repo_url} if repo_url else None
        return self.indexer.query_similar_chunks(query, k=10, filter_dict=filter_dict)
    
    def _extract_function_name(self, query: str) -> Optional[str]:
        """Extract function name from query."""
        import re
        
        # Common patterns
        patterns = [
            r"function\s+(\w+)",
            r"(\w+)\s+function",
            r"explain\s+(\w+)",
            r"what does\s+(\w+)",
            r"describe\s+(\w+)"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, query, re.IGNORECASE)
            if match:
                return match.group(1)
        
        return None
    
    def _extract_module_name(self, query: str) -> Optional[str]:
        """Extract module name from query."""
        import re
        
        patterns = [
            r"module\s+(\w+)",
            r"import\s+(\w+)",
            r"usage\s+of\s+(\w+)",
            r"use\s+(\w+)"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, query, re.IGNORECASE)
            if match:
                return match.group(1)
        
        return None
    
    def _extract_error_type(self, query: str) -> Optional[str]:
        """Extract error type from query."""
        import re
        
        patterns = [
            r"(\w+)\s+error",
            r"(\w+)\s+exception",
            r"handling\s+(\w+)"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, query, re.IGNORECASE)
            if match:
                return match.group(1)
        
        return None
    
    def _generate_response(self, query: str, query_type: str, context_docs: List[Document]) -> str:
        """Generate response based on query type and context."""
        
        # Format context
        context_str = self._format_context(context_docs)
        
        if query_type == 'function_explanation':
            function_name = self._extract_function_name(query)
            if function_name and context_docs:
                code = context_docs[0].page_content
                prompt = self.function_explanation_prompt
                response = prompt | self.llm | StrOutputParser()
                return response.invoke({
                    "code": code,
                    "context": context_str,
                    "function_name": function_name
                })
        
        elif query_type == 'module_usage':
            module_name = self._extract_module_name(query)
            if module_name:
                prompt = self.module_usage_prompt
                response = prompt | self.llm | StrOutputParser()
                return response.invoke({
                    "context": context_str,
                    "module_name": module_name
                })
        
        elif query_type == 'error_handling':
            error_type = self._extract_error_type(query)
            prompt = self.error_handling_prompt
            response = prompt | self.llm | StrOutputParser()
            return response.invoke({
                "context": context_str,
                "error_type": error_type or "general"
            })
        
        # Default general response
        prompt = self.general_prompt
        response = prompt | self.llm | StrOutputParser()
        return response.invoke({
            "context": context_str,
            "question": query
        })
    
    def _format_context(self, docs: List[Document]) -> str:
        """Format context documents for the LLM."""
        context_parts = []
        
        for doc in docs[:5]:  # Limit to top 5 documents
            metadata = doc.metadata
            
            header = f"File: {metadata.get('file_path', 'unknown')}"
            if 'start_line' in metadata:
                header += f" (lines {metadata['start_line']}-{metadata['end_line']})"
            
            context_parts.append(f"\n{header}\n```\n{doc.page_content}\n```")
        
        return "\n".join(context_parts)
    
    def explain_function(self, function_name: str, repo_url: str = None) -> str:
        """Explain a specific function."""
        return self.query_repository(f"explain the function {function_name}", repo_url)
    
    def find_module_usage(self, module_name: str, repo_url: str = None) -> str:
        """Find usage of a specific module."""
        return self.query_repository(f"how is the module {module_name} used", repo_url)
    
    def find_error_examples(self, error_type: str = "", repo_url: str = None) -> str:
        """Find error handling examples."""
        if error_type:
            query = f"find examples of handling {error_type} errors"
        else:
            query = "find error handling examples"
        
        return self.query_repository(query, repo_url)
    
    def get_repositories(self) -> List[Dict[str, Any]]:
        """Get list of indexed repositories."""
        # This would need to be implemented based on how we track repositories
        # For now, return empty list
        return []
    
    def clear_repository(self, repo_url: str) -> bool:
        """Clear all indexed data for a repository."""
        return self.indexer.remove_repository(repo_url)